from pydantic import BaseModel, Field, field_validator
from typing import Optional, Literal
from enum import Enum
from datetime import datetime
from dataclasses import dataclass

class TaskStatus(str, Enum):
    """任务状态枚举"""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

class FaceSwapRequest(BaseModel):
    """换脸请求模型 - 简化版（只支持extreme_beauty模式）"""
    source_image: str = Field(..., description="源图片base64编码，格式：")
    target_image: str = Field(..., description="目标图片base64编码，格式：")
    output_quality: int = Field(default=100, ge=1, le=100, description="输出图片质量(1-100)")
    
    @field_validator('source_image', 'target_image')
    def validate_base64_image(cls, v):
        if not v.startswith('data:image/'):
            raise ValueError('图片必须是data:image/格式的base64编码')
        return v

class FaceSwapResponse(BaseModel):
    """换脸响应模型"""
    task_id: str = Field(..., description="任务ID")
    message: str = Field(default="换脸任务已提交，请使用task_id查询进度")

class TaskQueryResponse(BaseModel):
    """任务查询响应模型"""
    task_id: str
    status: TaskStatus
    created_at: datetime
    updated_at: datetime
    result_url: Optional[str] = Field(None, description="结果图片CDN URL（任务完成后）")
    error_message: Optional[str] = Field(None, description="错误信息（任务失败时）")

class HealthResponse(BaseModel):
    """健康检查响应模型"""
    status: str
    gpu_available: bool
    gpu_memory_used: str
    gpu_memory_total: str
    active_tasks: int
    pending_tasks: int

@dataclass
class ProcessResult:
    """换脸处理结果数据类"""
    success: bool
    output_path: Optional[str] = None
    error_message: Optional[str] = None
